"""
Author: wenbin
Create: 2024-08-23
"""
import uuid
from dataclasses import dataclass
from operator import itemgetter
from typing import Dict, Any

from injector import inject
from langchain.memory import ConversationBufferWindowMemory
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.memory import BaseMemory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableConfig
from langchain_openai import ChatOpenAI
from langsmith.schemas import Run

import app.setup as setup
from app.exception import BusinessException
from app.exception.error import AppError
from app.models import App
from app.models.base_model import SQLAlchemy
from app.utils import format_documents
from app.validator.valid_app import ValidCreateApp
from .vector_database_service import VectorDatabaseService


@inject
@dataclass
class AppService:
    """
    应用业务层
    """
    db: SQLAlchemy
    vector_database_service: VectorDatabaseService

    def create(self, param: ValidCreateApp) -> None:
        """创建应用"""
        with self.db.auto_commit():
            app = App()
            app.account_id = uuid.uuid4()
            app.name = param.name.data
            app.icon = param.icon.data
            app.description = param.description.data
            self.db.session.add(app)

    def get(self, app_id: uuid) -> App:
        """查询应用"""
        app = self.db.session.query(App).filter_by(id=app_id).one()
        if app is None:
            raise BusinessException(AppError.APP_NOT_EXISTS)
        return app

    def debug(self, app_id: uuid, question: str) -> str:
        """应用对话调试"""
        # 构建提示词
        prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个强大的聊天机器人，能根据对应的上下文和历史对话信息正确回答用户的的问题"),
            ("system", "<context>{context}</context>"),
            MessagesPlaceholder("history"),
            ("human", "{question}")
        ])
        # 记忆组件
        memory = ConversationBufferWindowMemory(
            k=10,
            return_messages=True,
            chat_memory=RedisChatMessageHistory("abc", setup.REDIS_URL)
        )
        # 大模型
        llm = ChatOpenAI(model="gpt-3.5-turbo")

        # 构建链
        chain = (RunnablePassthrough.assign(
            history=RunnableLambda(_load_memory_variables) | itemgetter("history"),
            context=itemgetter(
                "question") | self.vector_database_service.get_retriever() | format_documents
        ) | prompt | llm | StrOutputParser()).with_listeners(on_end=_save_context)

        chain_input = {"question": question}
        content = chain.invoke(chain_input, config={
            "configurable": {
                "memory": memory
            }
        })
        return content


def _load_memory_variables(user_input: Dict[str, Any], config: RunnableConfig) -> Dict[str, Any]:
    """加载记忆变量"""
    configurable = config.get("configurable", {})
    memory = configurable.get("memory", None)
    if memory is not None and isinstance(memory, BaseMemory):
        return memory.load_memory_variables(user_input)
    return {"history": []}


def _save_context(run_obj: Run, config: RunnableConfig) -> None:
    """存储上下文信息到记忆组件"""
    configurable = config.get("configurable", {})
    memory = configurable.get("memory", None)
    if memory is not None and isinstance(memory, BaseMemory):
        memory.save_context(run_obj.inputs, run_obj.outputs)
